# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is dual-licensed under either the MIT license found in the
# LICENSE-MIT file in the root directory of this source tree or the Apache
# License, Version 2.0 found in the LICENSE-APACHE file in the root directory
# of this source tree. You may select, at your option, one of the
# above-listed licenses.

RustSources = provider(fields = {})

def srcs_arg():
    return {
        "srcs": attrs.list(
            attrs.source(),
            default = [],
            doc = """
    The set of Rust source files to be compiled by this target.

    One of the source files is the root module of the crate. By default this is
    `lib.rs` for libraries, `main.rs` for executables, or the crate's name with
    `.rs` appended. This can be overridden with the `crate_root` attribute.
""",
        ),
    }

def mapped_srcs_arg():
    return {
        "mapped_srcs": attrs.dict(
            key = attrs.source(),
            value = attrs.string(),
            sorted = False,
            default = {},
            doc = """
    Add source files along with a local path mapping. Rust is sensitive to the
    layout of source files, as the directory structure follows the module
    structure. However this is awkward if the source file is, for example,
    generated by another rule. In this case, you can set up a mapping from the
    actual source path to something that makes sense locally. For example
    `mapped_srcs = {":generate-module": "src/generated.rs"}`. These are added to
    the regular `srcs`, so a file should not be listed in both.
""",
        ),
    }

def srcs_filegroup_arg():
    return {
        "srcs_filegroup": attrs.option(
            attrs.dep(providers = [RustSources]),
            default = None,
            doc = """
    Directory of sources already combined together from `srcs` + `mapped_srcs`.
    This allows sources to be located in a different Buck package than the one
    containing the rust_library target which copmiles those sources.

    When using `srcs_filegroup`, the attributes `srcs` and `mapped_srcs` cannot
    also be passed, and `crate_root` must be passed.
""",
        ),
    }

def symlinked_srcs(ctx: AnalysisContext) -> Artifact:
    srcs_filegroup = getattr(ctx.attrs, "srcs_filegroup", None)
    if srcs_filegroup:
        if ctx.attrs.srcs:
            fail("srcs and srcs_filegroup cannot be used together")
        if ctx.attrs.mapped_srcs:
            fail("mapped_srcs and srcs_filegroup cannot be used together")
        if not ctx.attrs.crate_root:
            fail("crate_root is required when using srcs_filegroup")
        return srcs_filegroup[DefaultInfo].default_outputs[0]

    if not ctx.attrs.srcs and not ctx.attrs.mapped_srcs:
        fail("crate has no srcs, mapped_srcs, or srcs_filegroup")

    srcs = {src.short_path: src for src in ctx.attrs.srcs}
    srcs.update({k: v for v, k in ctx.attrs.mapped_srcs.items()})

    if "generated" in getattr(ctx.attrs, "labels", []):
        # For generated code targets, we always want to copy files in the [sources]
        # subtarget, never symlink.
        #
        # This ensures that IDEs that open the generated file always see the correct
        # directory structure.
        #
        # VS Code will expand symlinks when doing go-to-definition. In normal source
        # files this takes us back to the correct path, but for generated files the
        # expanded path may not be a well-formed crate layout.
        return ctx.actions.copied_dir("__srcs", srcs)
    else:
        # Decide whether to use symlinked_dir or copied_dir.
        #
        # If a source is a prefix of any other source, use copied_dir. This supports
        # e.g. `srcs = [":foo.crate"]` where :foo.crate is an http_archive, together
        # with a `mapped_srcs` which overlays additional generated files into that
        # directory. Symlinked_dir would error in this situation.
        prefixes = {}
        for src in sorted(srcs.keys(), key = len, reverse = True):
            if src in prefixes:
                return ctx.actions.copied_dir("__srcs", srcs)
            components = src.split("/")
            for i in range(1, len(components)):
                prefixes["/".join(components[:i])] = None

    # Otherwise, symlink it.
    return ctx.actions.symlinked_dir("__srcs", srcs)

def _rust_filegroup_impl(ctx: AnalysisContext) -> list[Provider]:
    return [
        DefaultInfo(default_output = symlinked_srcs(ctx)),
        RustSources(),
    ]

rust_filegroup = rule(
    impl = _rust_filegroup_impl,
    attrs = srcs_arg() | mapped_srcs_arg(),
)
